import argparse
import itertools
import inspect
import os
import numpy as np
import torch
import triton
import triton.language as tl
import triton.language.extra.libdevice as tldevice

FLA_SINGLE_THREAD = os.getenv("TRITON_PRECISION_NPU_SINGLE_THREAD", '0') == '1'

def single_wrapper(grid, grid_size, kernel):   
    if not FLA_SINGLE_THREAD:
        kwargs = {'USE_PIDS': False}
        for dim in range(grid_size):
            kwargs[f"PID{dim}"] = 0
        kernel(grid, kwargs)
        return
    assert isinstance(grid, tuple)
    
    grid_ranges = []
    for dim in grid:
        grid_ranges.append(range(dim))
    single_grid = (1,) * grid_size

    for combination in itertools.product(*grid_ranges):
        kwargs = {'USE_PIDS': True}
        for idx, dim in enumerate(combination):
            kwargs[f"PID{idx}"] = dim
        kernel(single_grid, kwargs)

supports_autotune_cache = "cache_results" in inspect.signature(triton.autotune).parameters
autotune_cache_kwargs = {"cache_results": FLA_CACHE_RESULTS} if supports_autotune_cache else {}

if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
    exp = tldevice.fast_expf
    exp2 = tldevice.exp2
    log = tldevice.fast_logf
    log2 = tldevice.fast_log2f
else:
    exp = tl.exp
    exp2 = tl.math.exp2
    log = tl.log
    log2 = tl.log2

@triton.heuristics({
    'HAS_DRESIDUAL': lambda args: args['dresidual'] is not None,
    'HAS_WEIGHT': lambda args: args['w'] is not None,
    'HAS_BIAS': lambda args: args['b'] is not None,
    'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None
})
@triton.autotune(
    configs=[
        triton.Config({}, num_warps=num_warps)
        for num_warps in [2, 4, 8, 16]
    ],
    key=['D', 'IS_RMS_NORM', 'STORE_DRESIDUAL', 'HAS_DRESIDUAL', 'HAS_WEIGHT'],
    **autotune_cache_kwargs
)
@triton.jit(do_not_specialize=['PID0'])
def layer_norm_gated_bwd_kernel1(
    x,  # pointer to the input
    g,  # pointer to the gate
    w,  # pointer to the weights
    b,  # pointer to the biases
    y,  # pointer to the output to be recomputed
    dy,  # pointer to the output gradient
    dx,  # pointer to the input gradient
    dg,  # pointer to the gate gradient
    dw,  # pointer to the partial sum of weights gradient
    db,  # pointer to the partial sum of biases gradient
    dresidual,
    dresidual_in,
    mean,
    rstd,
    T,
    BS,
    D: tl.constexpr,
    BD: tl.constexpr,
    ACTIVATION: tl.constexpr,
    IS_RMS_NORM: tl.constexpr,
    STORE_DRESIDUAL: tl.constexpr,
    HAS_DRESIDUAL: tl.constexpr,
    HAS_WEIGHT: tl.constexpr,
    HAS_BIAS: tl.constexpr,
    RECOMPUTE_OUTPUT: tl.constexpr,
    USE_PIDS: tl.constexpr,
    PID0
):
    if USE_PIDS:
        i_s = PID0
    else:
        i_s = tl.program_id(0)
    o_d = tl.arange(0, BD)
    mask = o_d < D
    x += i_s * BS * D
    g += i_s * BS * D
    if HAS_DRESIDUAL:
        dresidual += i_s * BS * D
    if STORE_DRESIDUAL:
        dresidual_in += i_s * BS * D
    dy += i_s * BS * D
    dx += i_s * BS * D
    dg += i_s * BS * D
    if RECOMPUTE_OUTPUT:
        y += i_s * BS * D
    if HAS_WEIGHT:
        b_w = tl.load(w + o_d, mask=mask).to(tl.float32)
        b_dw = tl.zeros((BD,), dtype=tl.float32)
    if HAS_BIAS:
        b_b = tl.load(b + o_d, mask=mask, other=0.0).to(tl.float32)
        b_db = tl.zeros((BD,), dtype=tl.float32)

    for i_t in range(i_s * BS, min(i_s * BS + BS, T)):
        # Load data to SRAM
        b_x = tl.load(x + o_d, mask=mask, other=0).to(tl.float32)
        b_g = tl.load(g + o_d, mask=mask, other=0).to(tl.float32)
        b_dy = tl.load(dy + o_d, mask=mask, other=0).to(tl.float32)

        if not IS_RMS_NORM:
            b_mean = tl.load(mean + i_t)
        b_rstd = tl.load(rstd + i_t)
        # Compute dx
        b_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
        b_xhat = tl.where(mask, b_xhat, 0.0)

        b_y = b_xhat * b_w if HAS_WEIGHT else b_xhat
        if HAS_BIAS:
            b_y = b_y + b_b
        if RECOMPUTE_OUTPUT:
            tl.store(y + o_d, b_y, mask=mask)

        b_sigmoid_g = tl.sigmoid(b_g)
        if ACTIVATION == 'swish':
            b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
            b_dy = b_dy * b_g * b_sigmoid_g
        elif ACTIVATION == 'silu':
            b_dg = b_dy * b_y * (b_sigmoid_g + b_g * b_sigmoid_g * (1 - b_sigmoid_g))
            b_dy = b_dy * b_g * b_sigmoid_g
        elif ACTIVATION == 'sigmoid':
            b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g)
            b_dy = b_dy * b_sigmoid_g
        b_wdy = b_dy
        if HAS_WEIGHT:
            b_wdy = b_dy * b_w
            b_dw += b_dy * b_xhat
        if HAS_BIAS:
            b_db += b_dy
        if not IS_RMS_NORM:
            b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D
            b_c2 = tl.sum(b_wdy, axis=0) / D
            b_dx = (b_wdy - (b_xhat * b_c1 + b_c2)) * b_rstd
        else:
            b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D
            b_dx = (b_wdy - b_xhat * b_c1) * b_rstd
        if HAS_DRESIDUAL:
            b_dres = tl.load(dresidual + o_d, mask=mask, other=0).to(tl.float32)
            b_dx += b_dres
        # Write dx
        if STORE_DRESIDUAL:
            tl.store(dresidual_in + o_d, b_dx, mask=mask)
        tl.store(dx + o_d, b_dx, mask=mask)
        tl.store(dg + o_d, b_dg, mask=mask)

        x += D
        g += D
        if HAS_DRESIDUAL:
            dresidual += D
        if STORE_DRESIDUAL:
            dresidual_in += D
        if RECOMPUTE_OUTPUT:
            y += D
        dy += D
        dx += D
        dg += D
    if HAS_WEIGHT:
        tl.store(dw + i_s * D + o_d, b_dw, mask=mask)
    if HAS_BIAS:
        tl.store(db + i_s * D + o_d, b_db, mask=mask)


def test_layer_norm_gated_bwd_kernel1(output_file):
    # 设置随机种子以确保可重复性
    torch.manual_seed(42)

    # 定义参数
    T = 64      # 序列长度
    D = 512     # 特征维度
    NS = 4      # 流处理器数量（模拟值）
    BS = 16     # 每个处理器的序列块大小
    BD = 512    # 特征维度的块大小（需要等于或大于D）

    # 生成随机输入张量
    device = 'npu'  # 确保使用NPU
    dtype = torch.float16

    # 输入张量
    x = torch.randn(T, D, dtype=dtype, device=device)
    g = torch.randn(T, D, dtype=dtype, device=device)
    weight = torch.randn(D, dtype=dtype, device=device)
    bias = torch.randn(D, dtype=dtype, device=device)
    dy = torch.randn(T, D, dtype=dtype, device=device)
    
    # 前向传播的统计量（均值和标准差）
    mean = torch.randn(T, dtype=dtype, device=device)
    rstd = torch.randn(T, dtype=dtype, device=device)
    
    # 可选的残差梯度
    dresidual = torch.randn(T, D, dtype=dtype, device=device)

    # 输出张量
    dx = torch.empty_like(x)
    dg = torch.empty_like(g)
    dresidual_in = torch.empty_like(x)
    y = torch.empty(T, D, dtype=dy.dtype, device=dy.device)  # 用于重新计算输出
    
    # 梯度累加张量
    dw = torch.empty((NS, D), dtype=torch.float32, device=weight.device)
    db = torch.empty((NS, D), dtype=torch.float32, device=bias.device)

    # 计算其他参数
    ACTIVATION = 'swish'
    IS_RMS_NORM = False
    STORE_DRESIDUAL = True
    HAS_DRESIDUAL = True
    HAS_WEIGHT = True
    HAS_BIAS = True
    RECOMPUTE_OUTPUT = True

    # 网格大小 - 一维网格，大小为NS
    grid = (NS,)

    # 调用内核函数
    def kernel(cur_grid, pid_args):
        layer_norm_gated_bwd_kernel1[cur_grid](
            x=x,
            g=g,
            w=weight,
            b=bias,
            y=y,
            dy=dy,
            dx=dx,
            dg=dg,
            dw=dw,
            db=db,
            dresidual=dresidual,
            dresidual_in=dresidual_in,
            mean=mean,
            rstd=rstd,
            T=T,
            BS=BS,
            D=D,
            BD=BD,
            ACTIVATION=ACTIVATION,
            IS_RMS_NORM=IS_RMS_NORM,
            STORE_DRESIDUAL=STORE_DRESIDUAL,
            HAS_DRESIDUAL=HAS_DRESIDUAL,
            HAS_WEIGHT=HAS_WEIGHT,
            HAS_BIAS=HAS_BIAS,
            RECOMPUTE_OUTPUT=RECOMPUTE_OUTPUT,
            **pid_args
        )
    single_wrapper(grid, 1, kernel)

    dg_numpy = dg.cpu().detach().numpy()
    np.savetxt(output_file, dg_numpy.reshape(-1, dg_numpy.shape[-1]))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Test Layer Norm Gated Backward Kernel 1d')
    parser.add_argument('--output', type=str, default='default_output.txt', 
                        help='Output file name (default: default_output.txt)')
    args = parser.parse_args()
    test_layer_norm_gated_bwd_kernel1(args.output)
